from copy import deepcopy
import numpy as np
import torch
from torch.optim import Adam
import gym
import time
import spinup.algos.pytorch.maddpg.core as core
from spinup.utils.logx import EpochLogger


class ReplayBuffer:
    """
    A simple FIFO experience replay buffer for DDPG agents.
    """

    def __init__(self, obs1_dim, act1_dim, obs2_dim, act2_dim, obs_dim, act_dim, size):
        self.obs1_buf = np.zeros(core.combined_shape(size, obs1_dim), dtype=np.float32)
        self.obs2_buf = np.zeros(core.combined_shape(size, obs2_dim), dtype=np.float32)
        self.obs_buf = np.zeros(core.combined_shape(size, obs_dim), dtype=np.float32)

        self.act1_buf = np.zeros(core.combined_shape(size, act1_dim), dtype=np.float32)
        self.act2_buf = np.zeros(core.combined_shape(size, act2_dim), dtype=np.float32)
        self.act_buf = np.zeros(core.combined_shape(size, act_dim), dtype=np.float32)

        self.rew1_buf = np.zeros(size, dtype=np.float32)
        self.rew2_buf = np.zeros(size, dtype=np.float32)

        self.next_obs1_buf = np.zeros(core.combined_shape(size, obs1_dim), dtype=np.float32)
        self.next_obs2_buf = np.zeros(core.combined_shape(size, obs2_dim), dtype=np.float32)
        self.next_obs_buf = np.zeros(core.combined_shape(size, obs_dim), dtype=np.float32)

        self.done_buf = np.zeros(size, dtype=np.float32)
        self.ptr, self.size, self.max_size = 0, 0, size

    def store(self, obs1, act1, obs2, act2, obs, act,
              rew1, rew2, next_obs1, next_obs2, next_obs, done):
        self.obs1_buf[self.ptr] = obs1
        self.obs2_buf[self.ptr] = obs2
        self.obs_buf[self.ptr] = obs

        self.act1_buf[self.ptr] = act1
        self.act2_buf[self.ptr] = act2
        self.act_buf[self.ptr] = act

        self.rew1_buf[self.ptr] = rew1
        self.rew2_buf[self.ptr] = rew2

        self.next_obs1_buf[self.ptr] = next_obs1
        self.next_obs2_buf[self.ptr] = next_obs2
        self.next_obs_buf[self.ptr] = next_obs

        self.done_buf[self.ptr] = done
        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)

    def sample_batch(self, batch_size=32):
        idxs = np.random.randint(0, self.size, size=batch_size)
        batch = dict(obs1=self.obs1_buf[idxs],
                     obs2=self.obs2_buf[idxs],
                     obs=self.obs_buf[idxs],

                     act1=self.act1_buf[idxs],
                     act2=self.act2_buf[idxs],
                     act=self.act_buf[idxs],

                     rew1=self.rew1_buf[idxs],
                     rew2=self.rew2_buf[idxs],

                     next_obs1=self.next_obs1_buf[idxs],
                     next_obs2=self.next_obs2_buf[idxs],
                     next_obs=self.next_obs_buf[idxs],

                     done=self.done_buf[idxs])
        return {k: torch.as_tensor(v, dtype=torch.float32) for k, v in batch.items()}


def maddpg(env_fn, actor_critic=core.MLPActorCritic, ac_kwargs=dict(), seed=0,
         steps_per_epoch=4000, epochs=100, replay_size=int(1e6), gamma=0.99,
         polyak=0.995, pi_lr=1e-3, q_lr=1e-3, batch_size=100, start_steps=15000000,
         update_after=1000, update_every=50, act_noise=0.1, num_test_episodes=10,
         max_ep_len=1000, logger_kwargs=dict(), save_freq=1):

    logger = EpochLogger(**logger_kwargs)
    logger.save_config(locals())

    torch.manual_seed(seed)
    np.random.seed(seed)

    env, test_env = env_fn(), env_fn()

    # agent1
    obs1_dim = env.observation_space1.shape
    act1_dim = env.action_space1.shape[0]
    act1_limit = env.action_space1.high[0]

    # agent2
    obs2_dim = env.observation_space2.shape
    act2_dim = env.action_space2.shape[0]
    act2_limit = env.action_space2.high[0]

    # all_agent
    obs_dim = env.observation_space.shape
    act_dim = env.action_space.shape[0]
    act_limit = env.action_space.high[0]


    # Create actor-critic module and target networks for agent1
    ac1 = actor_critic(env.observation_space1, env.action_space1,
                       env.observation_space, env.action_space, **ac_kwargs)
    ac1_targ = deepcopy(ac1)

    # Freeze target networks with respect to optimizers (only update via polyak averaging)
    for p1 in ac1_targ.parameters():
        p1.requires_grad = False

    # Create actor-critic module and target networks for agent2
    ac2 = actor_critic(env.observation_space2, env.action_space2,
                       env.observation_space, env.action_space, **ac_kwargs)
    ac2_targ = deepcopy(ac2)

    # Freeze target networks with respect to optimizers (only update via polyak averaging)
    for p2 in ac2_targ.parameters():
        p2.requires_grad = False

    # Experience buffer
    replay_buffer = ReplayBuffer(obs1_dim=obs1_dim, act1_dim=act1_dim,
                                 obs2_dim=obs2_dim, act2_dim=act2_dim,
                                 obs_dim=obs_dim, act_dim=act_dim,
                                 size=replay_size)

    # Count variables (protip: try to get a feel for how different size networks behave!)
    var_counts = tuple(core.count_vars(module) for module in [ac1.pi, ac1.q, ac2.pi, ac2.q])
    logger.log('\nNumber of parameters: \t ac1_pi: %d, \t ac1_q: %d, \t ac2_pi: %d, \t ac2_q: %d\n' % var_counts)

    # Set up function for computing agent1 Q-loss
    def compute_loss_q1(data):
        o = data['obs']
        a = data['act']
        n_o = data['next_obs']
        n_o1 = data['next_obs1']
        n_o2 = data['next_obs2']
        r1 = data['rew1']
        d = data['done']

        q = ac1.q(o, a)
        # Bellman backup for Q function
        with torch.no_grad():
            n_a1 = ac1_targ.pi(n_o1)
            n_a2 = ac2_targ.pi(n_o2)
            n_a = torch.cat([n_a1, n_a2], dim=1)
            q_pi1_targ = ac1_targ.q(n_o, n_a)
            backup = r1 + gamma * (1 - d) * q_pi1_targ
        # MSE loss against Bellman backup
        loss_q1 = ((q - backup) ** 2).mean()
        # Useful info for logging
        loss_info1 = dict(QVals=q.detach().numpy())

        return loss_q1, loss_info1

    # Set up function for computing agent1 pi loss
    def compute_loss_pi1(data):
        o1 = data['obs1']
        o2 = data['obs2']
        o = data['obs']

        a1 = ac1.pi(o1)
        a2 = ac2.pi(o2)
        a = torch.cat([a1, a2], dim=1)
        q_pi = ac1.q(o, a)
        return -q_pi.mean()

    # Set up optimizers for agent1 policy and q-function
    pi1_optimizer = Adam(ac1.pi.parameters(), lr=pi_lr)
    q1_optimizer = Adam(ac1.q.parameters(), lr=q_lr)

    # Set up model saving
    logger.setup_pytorch_saver1(ac1)

    def update1(data):
        # First run one gradient descent step for Q.
        q1_optimizer.zero_grad()
        loss_q1, loss_info1 = compute_loss_q1(data)
        loss_q1.backward()
        q1_optimizer.step()

        # Freeze Q-network so you don't waste computational effort
        # computing gradients for it during the policy learning step.
        for p1 in ac1.q.parameters():
            p1.requires_grad = False

        # Next run one gradient descent step for pi.
        pi1_optimizer.zero_grad()
        loss_pi1 = compute_loss_pi1(data)
        loss_pi1.backward()
        pi1_optimizer.step()

        # Unfreeze Q-network so you can optimize it at next DDPG step.
        for p1 in ac1.q.parameters():
            p1.requires_grad = True

        # Record things
        logger.store(LossQ1=loss_q1.item(), LossPi1=loss_pi1.item(), **loss_info1)

        # Finally, update target networks by polyak averaging.
        with torch.no_grad():
            for p1, p1_targ in zip(ac1.parameters(), ac1_targ.parameters()):
                # NB: We use an in-place operations "mul_", "add_" to update target
                # params, as opposed to "mul" and "add", which would make new tensors.
                p1_targ.data.mul_(polyak)
                p1_targ.data.add_((1 - polyak) * p1.data)

    # Set up function for computing agent2 Q-loss
    def compute_loss_q2(data):
        o = data['obs']
        a = data['act']
        n_o = data['next_obs']
        n_o1 = data['next_obs1']
        n_o2 = data['next_obs2']
        r2 = data['rew2']
        d = data['done']

        q = ac2.q(o, a)
        # Bellman backup for Q function
        with torch.no_grad():
            n_a1 = ac1_targ.pi(n_o1)
            n_a2 = ac2_targ.pi(n_o2)
            n_a = torch.cat([n_a1, n_a2], dim=1)
            q_pi2_targ = ac2_targ.q(n_o, n_a)
            backup = r2 + gamma * (1 - d) * q_pi2_targ
        # MSE loss against Bellman backup
        loss_q2 = ((q - backup) ** 2).mean()
        # Useful info for logging
        loss_info2 = dict(QVals=q.detach().numpy())

        return loss_q2, loss_info2

    # Set up function for computing agent1 pi loss
    def compute_loss_pi2(data):
        o1 = data['obs1']
        o2 = data['obs2']
        o = data['obs']

        a1 = ac1.pi(o1)
        a2 = ac2.pi(o2)
        a = torch.cat([a1, a2], dim=1)
        q_pi = ac1.q(o, a)
        return -q_pi.mean()

    # Set up optimizers for agent1 policy and q-function
    pi2_optimizer = Adam(ac2.pi.parameters(), lr=pi_lr)
    q2_optimizer = Adam(ac2.q.parameters(), lr=q_lr)

    # Set up model saving
    logger.setup_pytorch_saver2(ac2)

    def update2(data):
        # First run one gradient descent step for Q.
        q2_optimizer.zero_grad()
        loss_q2, loss_info2 = compute_loss_q2(data)
        loss_q2.backward()
        q2_optimizer.step()

        # Freeze Q-network so you don't waste computational effort
        # computing gradients for it during the policy learning step.
        for p2 in ac2.q.parameters():
            p2.requires_grad = False

        # Next run one gradient descent step for pi.
        pi2_optimizer.zero_grad()
        loss_pi2 = compute_loss_pi2(data)
        loss_pi2.backward()
        pi2_optimizer.step()

        # Unfreeze Q-network so you can optimize it at next DDPG step.
        for p2 in ac2.q.parameters():
            p2.requires_grad = True

        # Record things
        logger.store(LossQ2=loss_q2.item(), LossPi2=loss_pi2.item(), **loss_info2)

        # Finally, update target networks by polyak averaging.
        with torch.no_grad():
            for p2, p2_targ in zip(ac2.parameters(), ac2_targ.parameters()):
                # NB: We use an in-place operations "mul_", "add_" to update target
                # params, as opposed to "mul" and "add", which would make new tensors.
                p2_targ.data.mul_(polyak)
                p2_targ.data.add_((1 - polyak) * p2.data)

    def get_action1(o1, noise_scale):
        a = ac1.act(torch.as_tensor(o1, dtype=torch.float32))
        a += noise_scale * np.random.randn(act1_dim)
        return np.clip(a, -act_limit, act_limit)

    def get_action2(o2, noise_scale):
        a = ac1.act(torch.as_tensor(o2, dtype=torch.float32))
        a += noise_scale * np.random.randn(act2_dim)
        return np.clip(a, -act_limit, act_limit)

    def test_agent():
        for j in range(num_test_episodes):
            o1, o2, o, = test_env.reset()
            d, ep_ret1, ep_ret2, ep_len = False, 0, 0, 0
            while not (d or (ep_len == max_ep_len)):
                # Take deterministic actions at test time (noise_scale=0)
                o1, o2, o, r1, r2, d, _ = \
                    test_env.steps(get_action1(o1, 0), get_action2(o2, 0))
                ep_ret1 += r1
                ep_ret2 += r2
                ep_len += 1
            logger.store(TestEpRet1=ep_ret1, TestEpRet2=ep_ret2, TestEpLen=ep_len)

    # Prepare for interaction with environment
    total_steps = steps_per_epoch * epochs
    start_time = time.time()
    o1, o2, o,  = env.reset()
    d, ep_ret1, ep_ret2, ep_len = False, 0, 0, 0

    # Main loop: collect experience in env and update/log each epoch
    for t in range(total_steps):

        # Until start_steps have elapsed, randomly sample actions
        # from a uniform distribution for better exploration. Afterwards,
        # use the learned policy (with some noise, via act_noise).
        if t > start_steps:
            a1 = get_action1(o1, act_noise)
            a2 = get_action2(o2, act_noise)
            a = np.concatenate([a1, a2])
        else:
            a1 = env.action_space1.sample()
            a2 = env.action_space2.sample()
            a = np.concatenate([a1, a2])

        # Step the env
        o_n1, o_n2, o_n, r1, r2, d, _ = env.steps(a1, a2)
        ep_ret1 += r1
        ep_ret2 += r2
        ep_len += 1

        # Ignore the "done" signal if it comes from hitting the time
        # horizon (that is, when it's an artificial terminal signal
        # that isn't based on the agent's state)
        d = False if ep_len == max_ep_len else d

        # Store experience to replay buffer
        replay_buffer.store(o1, a1, o2, a2, o, a, r1, r2, o_n1, o_n2, o_n, d)

        # Super critical, easy to overlook step: make sure to update
        # most recent observation!
        o1 = o_n1
        o2 = o_n2
        o = o_n

        # End of trajectory handling
        if d or (ep_len == max_ep_len):
            logger.store(EpRet1=ep_ret1, EpRet2=ep_ret2, EpLen=ep_len)
            o1, o2, o, = env.reset()
            d, ep_ret1, ep_ret2, ep_len = False, 0, 0, 0
        # Update handling
        if t >= update_after and t % update_every == 0:
            for _ in range(update_every):
                batch = replay_buffer.sample_batch(batch_size)
                update1(data=batch)
                update2(data=batch)

        # End of epoch handling
        if (t + 1) % steps_per_epoch == 0:
            epoch = (t + 1) // steps_per_epoch

            # Save model
            if (epoch % save_freq == 0) or (epoch == epochs):
                logger.save_state({'env': env}, None)

            # Test the performance of the deterministic version of the agent.
            test_agent()

            # Log info about epoch
            logger.log_tabular('Epoch', epoch)
            logger.log_tabular('EpRet1', with_min_and_max=True)
            logger.log_tabular('EpRet2', with_min_and_max=True)
            logger.log_tabular('TestEpRet1', with_min_and_max=True)
            logger.log_tabular('TestEpRet2', with_min_and_max=True)
            logger.log_tabular('EpLen', average_only=True)
            logger.log_tabular('TestEpLen', average_only=True)
            logger.log_tabular('TotalEnvInteracts', t)
            logger.log_tabular('QVals', with_min_and_max=True)
            logger.log_tabular('LossPi1', average_only=True)
            logger.log_tabular('LossQ1', average_only=True)
            logger.log_tabular('LossPi2', average_only=True)
            logger.log_tabular('LossQ2', average_only=True)
            logger.log_tabular('Time', time.time() - start_time)
            logger.dump_tabular()


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--env', type=str, default='HalfCheetah-v2')
    parser.add_argument('--hid', type=int, default=256)
    parser.add_argument('--l', type=int, default=2)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--seed', '-s', type=int, default=0)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--exp_name', type=str, default='ddpg')
    args = parser.parse_args()

    from spinup.utils.run_utils import setup_logger_kwargs

    logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)

    maddpg(lambda: gym.make(args.env), actor_critic=core.MLPActorCritic,
         ac_kwargs=dict(hidden_sizes=[args.hid] * args.l),
         gamma=args.gamma, seed=args.seed, epochs=args.epochs,
         logger_kwargs=logger_kwargs)